from langgraph.checkpoint.memory import MemorySaver
from typing_extensions import TypedDict
from typing import Annotated
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, START, END
from utils.logger_config import LoggerSingleton
from tools_fc.Multiply import multiply_tool
from tools_fc.Add import add_tool
from llm_init.llm import llm
import json
from langchain_core.messages import ToolMessage
from langgraph.prebuilt import ToolNode, tools_condition
# 设置日志参数
logger = LoggerSingleton.get_logger()

# 定义状态
class State(TypedDict):
    messages: Annotated[list, add_messages]

# 初始化原始图
graph_builder = StateGraph(State)

# 添加工具节点
tools = [multiply_tool, add_tool]
tool_node = ToolNode(tools=tools)
graph_builder.add_node("tools", tool_node)

# 添加 LLM 模型节点，并绑定工具
llm_with_tools = llm.bind_tools(tools)

def chatbot(state: State):
    """调用 LLM 处理消息"""
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

graph_builder.add_node("chatbot", chatbot)

# 设定检查点存储
task_memory = MemorySaver()

# 添加边：控制对话流程
graph_builder.add_conditional_edges("chatbot", tools_condition)
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")

# 编译图
graph = graph_builder.compile(checkpointer=task_memory)



if __name__ == "__main__":
    # 绘图——可视化
    flag_visual = 0
    if flag_visual != 0:
        logger.info('graph可视化')
        from utils.visual_drawing import save_graph_image
        save_graph_image(graph)
    else:
        logger.info('graph未可视化')

    config = {"configurable": {"thread_id": "1"}}
    def stream_graph_updates(user_input: str):
        events = graph.stream(
            {"messages": [{"role": "user", "content": user_input}]},
            config,
            stream_mode="values",
        )
        for event in events:
            event["messages"][-1].pretty_print()

    while True:
        try:
            user_input = input("User(quit\exit\q 退出): ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("Goodbye!")
                break

            stream_graph_updates(user_input)
        except:
            # fallback if input() is not available
            user_input = "What do you know about LangGraph?"
            print("User: " + user_input)
            stream_graph_updates(user_input)
            break